import pymysql
import pymysql.cursors
from yaml import full_load
from kubernetes import client, config
from datetime import *
import json
import threading

CHAOS_MESH = {
    "host": "1.92.152.201",
    "user": "root",
    "password": "elastic",
    "port": 3306,
    "database": "chaos_mesh",
    "kubernetes_path": "/loganomaly/data/kubeconfig/config",
}


def get_stressor_type(stressors):
    failure_type = list(stressors.keys())[0]
    return failure_type


def get_service_for_pod(cmdb_id_list):
    service_list = [cmdb_id.split("-")[0] for cmdb_id in cmdb_id_list]
    return service_list


def str2int_time(duration_str):
    units = {"s": 1, "m": 60, "h": 3600, "d": 86400}

    if duration_str[-1] not in units:
        raise ValueError("Invalid time unit. Allowed units: 's', 'm', 'h', 'd'")

    try:
        seconds = int(duration_str[:-1]) * units[duration_str[-1]]
        return seconds
    except ValueError:
        raise ValueError(
            "Invalid duration format. Expected format: '<number><unit>', e.g. '15m'"
        )


class ChaosMeshConnection:
    def __init__(self, _config):
        self.connection = None
        try:
            self.connection = pymysql.connect(
                host=_config["host"],
                user=_config["user"],
                password=_config["password"],
                database=_config["database"],
                cursorclass=pymysql.cursors.DictCursor,
            )

        except pymysql.Error as error:
            print("Error connecting to MySQL database:", error)
        config.kube_config.load_kube_config(config_file=_config["kubernetes_path"])
        self.v1 = client.CoreV1Api()
        self._config = _config
        self.rlock = threading.Lock()

    def __del__(self):
        self.connection.close()

    def get_pods_for_label(self, namespace, label):
        label_selector = f"app={label}"
        pods = self.v1.list_namespaced_pod(namespace, label_selector=label_selector)
        pods = pods.items
        pods = [pod.metadata.name for pod in pods]
        return pods

    def get_culprit_cmdb_id(self, cmdb_id_list, *selectors):
        culprit_cmdb_id_list = []

        for selector in selectors:
            if "pods" in selector:
                pods = selector["pods"][selector["namespaces"][0]]
                culprit_cmdb_id_list.extend(
                    [pod for pod in pods if pod in cmdb_id_list]
                )
            else:
                namespace = selector["namespaces"][0]
                label = selector["labelSelectors"]["app"]
                pods = self.get_pods_for_label(namespace=namespace, label=label)
                culprit_cmdb_id_list.extend(pods)
        return culprit_cmdb_id_list

    def find_duration_in_workflow(self, metadata):
        workflow = metadata["labels"]["chaos-mesh.org/workflow"]
        instance = "-".join(
            metadata["labels"]["chaos-mesh.org/controlled-by"].split("-")[:-1]
        )
        # workflow
        sql = "SELECT * FROM workflow_entities WHERE name = %s"
        workflow_result = self.query(sql, (workflow,))
        workflow_result = workflow_result[0]["workflow"]
        workflow_data = json.loads(workflow_result)
        templates = workflow_data["spec"]["templates"]
        case = [case for case in templates if case["name"] == instance]
        return case[0]["deadline"]

    def feature_extract(self, experiment, cmdb_id_list, kind, metadata):
        service_list = []
        culprit_cmdb_id_list = []
        failure_type = ""
        duration = 0
        if "duration" in experiment.keys():
            duration = str2int_time(experiment["duration"])
        else:
            duration = str2int_time(self.find_duration_in_workflow(metadata))
        if kind == "NetworkChaos":
            failure_type = experiment["action"]
            direction = experiment["direction"]

            if direction == "to":
                selector = experiment["selector"]
                culprit_cmdb_id_list = self.get_culprit_cmdb_id(cmdb_id_list, selector)
            elif direction == "from":
                target = experiment["target"]["selector"]
                culprit_cmdb_id_list = self.get_culprit_cmdb_id(cmdb_id_list, target)
            else:
                selector = experiment["selector"]
                target = experiment["target"]["selector"]
                culprit_cmdb_id_list = self.get_culprit_cmdb_id(
                    cmdb_id_list, selector, target
                )
        elif kind == "PodChaos":
            failure_type = experiment["action"]
            selector = experiment["selector"]
            culprit_cmdb_id_list = self.get_culprit_cmdb_id(cmdb_id_list, selector)
        elif kind == "StressChaos":
            failure_type = get_stressor_type(experiment["stressors"])
            selector = experiment["selector"]
            culprit_cmdb_id_list = self.get_culprit_cmdb_id(cmdb_id_list, selector)
        service_list = get_service_for_pod(culprit_cmdb_id_list)
        return service_list, culprit_cmdb_id_list, failure_type, duration

    def group_by_object_id(self, start_date, end_date):
        sql = f"SELECT * FROM events WHERE created_at >= '{start_date}' AND created_at < '{end_date}' AND reason='Applied'"
        # query = f"SELECT * FROM events WHERE reason='Applied'"
        result = self.query(sql)

        # 创建一个空字典用于存储按object_id分组后的结果
        grouped_result = {}

        # 遍历result中的每行记录
        for row in result:
            object_id = row["object_id"]

            temp = {}
            temp["object_id"] = object_id
            # 转换created_at为timestamp
            temp["created_at"] = int(datetime.timestamp(row["created_at"]))
            temp["kind"] = row["kind"]
            temp["cmdb_id"] = []

            # 如果object_id不存在于grouped_result中，则创建一个新的列表，并将当前记录添加到列表中
            if object_id not in grouped_result:
                grouped_result[object_id] = temp

            # 处理message字段
            message = row["message"]
            message = message[len("Successfully apply chaos for ") :]
            cmdb_id = message.split("/")[1].strip()

            grouped_result[object_id]["cmdb_id"].append(cmdb_id)
        return grouped_result

    def query(self, sql, args=None, retry=0):
        self.rlock.acquire()
        self.connection.ping()
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(sql, args)
                result = cursor.fetchall()
        except Exception as e:
            self.rlock.release()
            if retry >= 5:
                raise Exception(repr(e))
            return self.query(sql, args, retry + 1)
        self.rlock.release()
        return result

    def create_ground_truth(self, start_time, end_time):
        start_date = datetime.fromtimestamp(start_time).replace(tzinfo=None)
        end_date = datetime.fromtimestamp(end_time).replace(tzinfo=None)
        grouped_result = self.group_by_object_id(start_date, end_date)
        ts_list = []
        service_list = []
        cmdb_id_list = []
        failure_type_list = []
        duration_list = []
        print(f"grouped_result length: {len(grouped_result)}")
        for object_id, row in grouped_result.items():
            ts = row["created_at"]
            # experiments
            sql = "SELECT * FROM experiments WHERE uid = %s"
            experiments_result = self.query(sql, (object_id,))

            if len(experiments_result) == 0:
                continue
            experiment_data = experiments_result[0][
                "experiment"
            ]  # 如果此处没有数据怎么办??
            experiment = json.loads(experiment_data)
            # print(f"experiment: {experiment}")

            try:
                culprit_service_list, culprit_cmdb_id_list, failure_type, duration = (
                    self.feature_extract(
                        experiment["spec"],
                        row["cmdb_id"],
                        row["kind"],
                        experiment["metadata"],
                    )
                )
            except Exception as e:
                print(repr(e))
                continue
            for service, cmdb_id in zip(culprit_service_list, culprit_cmdb_id_list):
                ts_list.append(ts)
                service_list.append(service)
                cmdb_id_list.append(cmdb_id)
                failure_type_list.append(failure_type)
                duration_list.append(duration)

        ground_truth = {
            "timestamp": ts_list,
            "service": service_list,
            "cmdb_id": cmdb_id_list,
            "failure_type": failure_type_list,
            "duration": duration_list,
        }

        return ground_truth
